from manim import *
from PIL import Image, ImageDraw
import imageio
import random
import numpy as np


# Part 1: 用Manim创建基础场景
class BouncingBall(Scene):
    def construct(self):
        # 创建小球（AI随机选择颜色）
        colors = [RED, GREEN, BLUE, YELLOW, PURPLE]
        ball = Circle(radius=0.5, color=random.choice(colors)).shift(UP)

        # 创建弹跳动画
        self.play(
            ball.animate.shift(DOWN * 2).set_color(random.choice(colors)),
            rate_func=lambda t: np.abs(np.sin(t * 3 * np.pi)),  # 弹跳效果
            run_time=2
        )
        self.wait(0.5)


# Part 2: 用Pillow生成AI增强的帧
def generate_ai_frame(frame_num):
    # 创建画布
    img = Image.new("RGB", (800, 600), "white")
    draw = ImageDraw.Draw(img)

    # AI生成随机参数（这里用随机数模拟）
    x = 400 + 300 * np.sin(frame_num * 0.2)
    y = 300 + 200 * np.cos(frame_num * 0.3)
    r = 20 + 15 * np.sin(frame_num * 0.4)
    color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

    # 绘制动态变化的圆
    draw.ellipse([x-r, y-r, x+r, y+r], fill=color)

    # 添加AI生成的文字
    draw.text((50, 50),
              f"AI Frame: {frame_num}",
              fill="black",
              font_size=30)


    return img


# 生成动画
def create_animation():
    # 生成Manim动画
    scene = BouncingBall()
    scene.render()

    # 生成Pillow AI帧动画
    frames = [] # 列表（list）：有序、可变
    for i in range(30):
        frame = generate_ai_frame(i)
        frames.append(np.array(frame))

    # 保存为GIF
    imageio.mimsave('ai_animation.gif', frames, duration=0.5)
    print("动画已保存为 ai_animation.gif")


if __name__ == "__main__":
    create_animation()
